#!/bin/bash/python
# _*_ coding: utf-8 _*_
# file： task_upgrade_util
# date-time: 2023/9/14-9:30
# auth:hard work
# description

import time
import os
from concurrent.futures import ThreadPoolExecutor

from tools_dev_and_sys.tool_ftp_get_put import put_local_file_to_os
from tools_dev_and_sys.tool_ssh_base import SSH_Server
import threading

def base_cfg(device_type):
    if device_type == "gw":
        version_file = "TSN_Release_v{}-usrfs.tar.bz2"
        kernel_file = "nari_v3.fit"
        local_dir = r'C:/version/gw/'
        return version_file,kernel_file,local_dir
    elif device_type == 'v5':
        version_file = 'kg6524.tar.bz2'
        kernel_file = 'hengshan_v5.fit'
        local_dir = r'C:/version/V5/'
        return version_file, kernel_file, local_dir

def kernel_auto_upgrade(ip, usrname, passwd, port):
    # _version = input("select version [date]: ")
    _file_name = "nari_v3.fit"
    local_dir = r"C:\\version\\GW\\" + _file_name
    remote_dir = '/tmp/' + _file_name
    client = SSH_Server(ip, usrname, passwd, port)
    try:
        client.connect()
        client.open_channel()
        if os.path.isfile(local_dir):
            print("found  file ", _file_name)
            client.channel_send_cmd(cmd='cd /tmp/', break_prompt='# ')
            res = client.channel_send_cmd(cmd='find -name {}'.format(_file_name), break_prompt='# ')
            if "nari_v3.fit" in res:
                print('file is exits,rename it to nari_v3.fit_old ..')
            client.channel_send_cmd(cmd='mv nari_v3.fit nari_v3.fit_old', break_prompt='# ')

            ret = put_local_file_to_os(client, local_dir, remote_dir, _file_name)
            if ret:
                print("upgrade start ,please wait")
                client.channel_send_cmd(cmd='chmod 755 %s' % _file_name, break_prompt='# ')
                client.channel_send_cmd(cmd='flashcp -v ./%s /dev/mtd0' % _file_name, break_prompt='# ')
                time.sleep(10)
            else:
                print('file not exits')
        else:
            raise "local file not found"

        client.channel_send_cmd(cmd='reboot', break_prompt='# ')
        print("system is reboot, wait ...")
    except FileNotFoundError as e:
        print("local file not found", e)
    except Exception as e:
        print(e)
    else:
        client.close()

def stack_auto_upgrade(ip, usrname, passwd, port):
    _version = '2.23'
    _file_name = "TSN_Release_v{}-usrfs.tar.bz2".format(_version)
    local_dir = r"C:\\version\\GW\\" + _file_name
    remote_dir = '/app/' + _file_name
    client = SSH_Server(ip, usrname, passwd, port)
    try:
        client.connect()
        client.open_channel()
        if os.path.isfile(local_dir):
            res = client.channel_send_cmd(cmd='find -name "usrfs*"', break_prompt='# ')
            if './usrfs' in res:
                print('find same file,rename it to usrfs_old')
                client.channel_send_cmd(cmd='mv usrfs usrfs_old', break_prompt='# ')
            ret = put_local_file_to_os(client, local_dir, remote_dir, _file_name)
            if ret:
                print("upgrade start ,please wait")
                client.channel_send_cmd(cmd='tar -xjvf %s' % _file_name, break_prompt='# ')
                client.channel_send_cmd(cmd='tar -xjvf usrfs.tar.bz2', break_prompt='# ')
                time.sleep(30)
                res = client.channel_send_cmd(cmd='find -name "usrfs*"', break_prompt='# ')
                if './usrfs' in res:
                    print('auto upgrade success')
                    client.channel_send_cmd(cmd='cd usrfs', break_prompt='# ')
                    client.channel_send_cmd(cmd='./gen_bringup.sh kgtsn gw', break_prompt='# ')
                else:
                    print('usrfs file not found')
            else:
                print('file not found')

        client.channel_send_cmd(cmd='reboot', break_prompt='# ')
        print("system is reboot, wait ...")
    except FileNotFoundError as e:
        print("local file not found", e)
    except Exception as e:
        print(e)
    else:
        client.close()

def main():
    usrname = "root"
    passwd = "hengshan"
    port = 22
    ip_lst = input("input dut ip, Separate by ',', max num 5. ?>>> ")
    ip_lst = ip_lst.split(',')

    t1 = threading.Thread(target=stack_auto_upgrade,args=(ip_lst[0],usrname,passwd,port))
    t2 = threading.Thread(target=stack_auto_upgrade,args=(ip_lst[1],usrname,passwd,port))
    t3 = threading.Thread(target=stack_auto_upgrade,args=(ip_lst[2],usrname,passwd,port))
    t4 = threading.Thread(target=stack_auto_upgrade,args=(ip_lst[3],usrname,passwd,port))
    t5 = threading.Thread(target=stack_auto_upgrade,args=(ip_lst[4],usrname,passwd,port))
    t1.start()
    t2.start()
    t3.start()
    t4.start()
    t5.start()

    pool = ThreadPoolExecutor(max_workers=len(ip_lst))
    pool.submit(stack_auto_upgrade,ip_lst,usrname,passwd,port)


if __name__ == '__main__':
    main()
